import os.path

import pandas as pd
import torch

import src.RL.TrainDQN
from src.RL.QNet import DQN
from src.RL.FinEnv import StockTradingEnvironment
from src.RL.TrainDQN import WINDOW_SIZE, state2tensor
from src.Utils.MyUtil import LinePlot
from Estimate import information_ratio, sharpe_ratio, max_drawdown
from src.EnvironmentVariables import BASE_PATH
import matplotlib.pyplot as plt
from src.RL.TCNQnet import QNetWithTCN
from src.RL.EasyNet import EasyQNet


class RLSingleTester:
    def __init__(self, stockdata: pd.DataFrame, base_model, titleName='', window_size=WINDOW_SIZE):
        self.dqn = DQN(base_model, QNetWork=EasyQNet)
        self.dqn.eval()
        self.stockdata = stockdata
        self.titleName = titleName
        self.date = stockdata[65 + window_size + 6:]
        self.window_size = window_size
        self.date = self.date['date']

    def test(self, input_col=None):
        with torch.no_grad():
            if input_col is None:
                input_col = StockTradingEnvironment.INPUT_COL
            kwargs4env = {
                'stock_data': self.stockdata,
                'window_size': self.window_size
            }
            env = StockTradingEnvironment(**kwargs4env)
            state = env.reset(**kwargs4env)
            total_wealth = [1]
            stock_wealth = [1]
            real_profit_list = [0]
            stock_profit_list = [0]
            reward_list = [0]
            while True:
                state = state2tensor(state, input_col)
                action = self.dqn.choose_action(state)
                next_state, reward, done = env.step(action)
                reward, real_profit, keep_profit = reward
                reward_list.append(reward / 100)
                real_profit_list.append(real_profit)
                stock_profit_list.append(keep_profit)
                next_state = state2tensor(next_state, input_col)

                total_wealth.append(total_wealth[-1] * (1 + real_profit))
                stock_wealth.append(stock_wealth[-1] * (1 + keep_profit))
                state = next_state
                if done:
                    break
            return total_wealth, stock_wealth, real_profit_list, stock_profit_list, reward_list, env.position_his

    def test_and_show(self, use_col=None):
        total_wealth, stock_wealth, real_profit_list, stock_profit_list, reward_list, action_list = self.test(
            use_col)
        ir = information_ratio(real_profit_list, stock_profit_list)
        sr = sharpe_ratio(real_profit_list, risk_free_rate=0.0001)
        stock_sr = sharpe_ratio(stock_profit_list, 0.0001)
        mdd = max_drawdown(stock_profit_list)
        stock_mdd = max_drawdown(stock_profit_list)
        plot = LinePlot(title=self.titleName)
        plot.draw(self.date,
                  stock_wealth,
                  total_wealth,
                  action_list,
                  reward_list,
                  real_profit_list,
                  stock_profit_list,
                  save=False,
                  close=False)
        print(f"IR:{ir:.4f};\n"
              f"Sharpe Ration Model:{sr:.4f} / Stock:{stock_sr:.4f}\n"
              f"Max Draw Down Model:{mdd:.4f} / Stock:{stock_mdd:.4f}\n")


def estimate(stock_id, splitTime='2010-1-1'):
    StockTradingEnvironment.INPUT_COL = EasyQNet.NEED_COL

    dir_path = os.path.join(BASE_PATH, 'data/preProcessedIndexData/sz50stockSignal')
    files_path = os.path.join(dir_path, stock_id + '.csv')
    df = pd.read_csv(files_path, parse_dates=['date'])
    df = df.sort_values('date').dropna()
    df = df[df['date'] >= pd.to_datetime(splitTime)]
    RLSingleTester(df, os.path.join(BASE_PATH, 'Models/RL_model/sh600519checkpoint10.pt'),
                   f'{stock_id}', window_size=1).test_and_show()

    plt.show()


if __name__ == '__main__':
    estimate('sh600519')
